/*
 *  +----------------------------------------------------------------------
 *  | sophon [ A FAST GAME FRAMEWORK ]
 *  +----------------------------------------------------------------------
 *  | Copyright (c) 2023-2029 All rights reserved.
 *  +----------------------------------------------------------------------
 *  | Licensed ( http:www.apache.org/licenses/LICENSE-2.0 )
 *  +----------------------------------------------------------------------
 *  | Author: jqiris <1920624985@qq.com>
 *  +----------------------------------------------------------------------
 */

use std::net::SocketAddr;
use std::sync::Arc;

use async_trait::async_trait;
use tokio::net::{TcpListener, TcpStream};
use tokio_rustls::{server::TlsStream, TlsAcceptor};
use tokio_tungstenite::WebSocketStream;

use crate::{error, info};
use crate::configs::SslConf;
use crate::encryption::make_server_config;
use crate::logger::*;

//websocket处理
#[async_trait]
pub trait IWebsocketHandler {
    async fn handle(&self, addr: SocketAddr, stream: WsStream);
}

pub type WebsocketHandler = Arc<dyn IWebsocketHandler + Send + Sync>;

#[derive(Debug)]
pub enum WsStream {
    Tcp(WebSocketStream<TcpStream>),
    Tls(WebSocketStream<TlsStream<TcpStream>>),
}

pub struct Websocket {
    host: String,
    port: i32,
    ssl: SslConf,
    handler: WebsocketHandler,
}

impl Websocket {
    pub fn new(host: String, port: i32, ssl: SslConf, handler: WebsocketHandler) -> Self {
        Self {
            host,
            port,
            ssl,
            handler,
        }
    }
    async fn listen_and_serve(&self) {
        let url = format!("{}:{}", self.host, self.port);
        let listener = TcpListener::bind(url).await.expect("failed to bind");
        while let Ok((stream, addr)) = listener.accept().await {
            let handler = self.handler.clone();
            tokio::spawn(Self::handle_connection(stream, addr, handler));
        }
    }

    async fn handle_connection(raw_stream: TcpStream, addr: SocketAddr, handler: WebsocketHandler) {
        info!("Incoming TCP connection from: {}", addr);
        let ws_stream = tokio_tungstenite::accept_async(raw_stream)
            .await
            .expect("Error during the websocket handshake occurred");
        info!("WebSocket connection established: {}", addr);
        handler.handle(addr, WsStream::Tcp(ws_stream)).await;
    }

    async fn listen_and_serve_tls(&self) {
        let config = make_server_config(&self.ssl.cert_file, &self.ssl.key_file).unwrap();
        let url = format!("{}:{}", self.host, self.port);
        let acceptor = TlsAcceptor::from(Arc::new(config));
        let listener = TcpListener::bind(url).await.expect("failed to bind");
        while let Ok((stream, addr)) = listener.accept().await {
            let accept_resp = acceptor.accept(stream).await;
            match accept_resp {
                Ok(stream) => {
                    let handler = self.handler.clone();
                    tokio::spawn(Self::handle_connection_tls(stream, addr, handler));
                }
                Err(err) => {
                    error!("tls connection err:{}", err);
                }
            }
        }
    }
    async fn handle_connection_tls(
        raw_stream: TlsStream<TcpStream>,
        addr: SocketAddr,
        handler: WebsocketHandler,
    ) {
        info!("Incoming TCP connection from: {}", addr);
        let ws_stream = tokio_tungstenite::accept_async(raw_stream)
            .await
            .expect("Error during the websocket handshake occurred");
        info!("WebSocket connection established: {}", addr);
        handler.handle(addr, WsStream::Tls(ws_stream)).await;
    }

    pub async fn run(&self) {
        if self.ssl.power_on {
            self.listen_and_serve_tls().await;
        } else {
            self.listen_and_serve().await;
        }
    }
}
